import numpy as np
import os
import output_function as of

# 使用示例
base_path = './data/SOL/'  # 根据您的保存路径进行修改

# 首先，只加载一次初始数据和模拟参数
n_init0_matrix, xx, yy, x, y = of.load_initial_data()
config = of.load_simulation_config()

# 从config字典中提取特定的参数并赋值给变量
dx = config['dx']
dy = config['dy']
Lx = config['Lx']
Ly = config['Ly']
Ly_in=config['Ly_in']
Ly_out=config['Ly_out']
Nx = config['Nx']
Ny = config['Ny']
n_up = config['n_up']
n_0 = config['n_0']
n_down = config['n_down']
Delta_n = config['Delta_n']
L_perp = config['L_perp']
L_paral = config['L_paral']
g_hat = config['g_hat']
zeta = config['zeta']
D_n = config['D_n']
mu = config['mu']
Ra_star = config['Ra_star']
Pr = config['Pr']
Omega_freq = config['Omega_freq']
alpha_sh0 = config['alpha_sh0']
dt = config['dt']
ntime = config['ntime']
ndiag = config['ndiag']

# 假设我们从out_iter=0开始，直到8000，且每次增加100
for out_iter in range(0, 9901, 100):
    # 加载数据
    n, psi, omega, u, v, out_iter_loaded = of.load_simulation_data(base_path, out_iter)
    
    # 确保加载了正确的时间步长
    if out_iter_loaded != out_iter:
        print(f"Failed to load the correct time step {out_iter} or file does not exist.")
        continue  # 如果加载失败，跳过当前迭代，继续下一个时间步
    
    # 打印一些变量以确认它们已被正确加载和赋值
    print(f"Successfully loaded data for time step: {out_iter_loaded}")
    
    # 使用加载的数据进行分析、绘图或其他处理
    of.diag_n_pcolor_filepath(n, n_init0_matrix, u, v, xx, yy, out_iter, Ra_star, dt,base_folder='./SOL_plot')
    of.plot_and_save_n_density_x_section(x,y, n, n_init0_matrix, Ny//4, out_iter)
    of.plot_and_save_v_x_section(x,y, v, Ny//4, out_iter)